#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <nettle/sm4.h>

#include "abcrypto.h"
#include "cipher.h"

#define BLOCK_SIZE 16

// ECB 模式分段加密
static void sm4_ecb_encrypt_segment(struct sm4_ctx *ctx,
				    const uint8_t *plaintext, size_t length,
				    uint8_t *ciphertext)
{
	// sm4_set_encrypt_key(ctx, key);
	for (size_t i = 0; i < length; i += BLOCK_SIZE) {
		sm4_crypt(ctx, BLOCK_SIZE, ciphertext + i, plaintext + i);
	}
}

// ECB 模式分段解密
static void sm4_ecb_decrypt_segment(struct sm4_ctx *ctx,
				    const uint8_t *ciphertext, size_t length,
				    uint8_t *plaintext)
{
	// sm4_set_decrypt_key(&ctx, key);
	for (size_t i = 0; i < length; i += BLOCK_SIZE) {
		sm4_crypt(ctx, BLOCK_SIZE, plaintext + i, ciphertext + i);
	}
}

// CBC 模式分段加密
static void sm4_cbc_encrypt_segment(struct sm4_ctx *ctx, uint8_t *iv,
				    const uint8_t *plaintext, size_t length,
				    uint8_t *ciphertext)
{
	// sm4_set_encrypt_key(&ctx, key);
	for (size_t i = 0; i < length; i += BLOCK_SIZE) {
		uint8_t block[BLOCK_SIZE];
		for (size_t j = 0; j < BLOCK_SIZE; j++) {
			block[j] = plaintext[i + j] ^ iv[j];
		}
		sm4_crypt(ctx, BLOCK_SIZE, ciphertext + i, block);
		memcpy(iv, ciphertext + i, BLOCK_SIZE);
	}
}

// CBC 模式分段解密
static void sm4_cbc_decrypt_segment(struct sm4_ctx *ctx, uint8_t *iv,
				    const uint8_t *ciphertext, size_t length,
				    uint8_t *plaintext)
{
	// sm4_set_decrypt_key(ctx, key);
	for (size_t i = 0; i < length; i += BLOCK_SIZE) {
		uint8_t block[BLOCK_SIZE];
		sm4_crypt(ctx, BLOCK_SIZE, block, ciphertext + i);
		for (size_t j = 0; j < BLOCK_SIZE; j++) {
			plaintext[i + j] = block[j] ^ iv[j];
		}
		memcpy(iv, ciphertext + i, BLOCK_SIZE);
	}
}

// CTR 模式分段加密/解密（CTR 模式加解密操作相同）
static void sm4_ctr_encrypt_decrypt_segment(struct sm4_ctx *ctx, uint8_t *ctr,
					    const uint8_t *input, size_t length,
					    uint8_t *output)
{
	// sm4_set_encrypt_key(&ctx, key);
	uint8_t keystream[BLOCK_SIZE];
	for (size_t i = 0; i < length; i += BLOCK_SIZE) {
		sm4_crypt(ctx, BLOCK_SIZE, keystream, ctr);
		for (size_t j = 0; j < BLOCK_SIZE && i + j < length; j++) {
			output[i + j] = input[i + j] ^ keystream[j];
		}
		// 增加计数器
		for (int j = BLOCK_SIZE - 1; j >= 0; j--) {
			if (++ctr[j] != 0) {
				break;
			}
		}
	}
}

// CFB 模式分段加密
static void sm4_cfb_encrypt_segment(struct sm4_ctx *ctx, uint8_t *iv,
				    const uint8_t *plaintext, size_t length,
				    uint8_t *ciphertext)
{
	// struct sm4_ctx ctx;
	// sm4_set_encrypt_key(&ctx, key);
	for (size_t i = 0; i < length; i += BLOCK_SIZE) {
		uint8_t block[BLOCK_SIZE];
		sm4_crypt(ctx, BLOCK_SIZE, block, iv);
		for (size_t j = 0; j < BLOCK_SIZE && i + j < length; j++) {
			ciphertext[i + j] = plaintext[i + j] ^ block[j];
		}
		memcpy(iv, ciphertext + i, BLOCK_SIZE);
	}
}

// CFB 模式分段解密
static void sm4_cfb_decrypt_segment(struct sm4_ctx *ctx, uint8_t *iv,
				    const uint8_t *ciphertext, size_t length,
				    uint8_t *plaintext)
{
	// sm4_set_encrypt_key(&ctx, key);
	for (size_t i = 0; i < length; i += BLOCK_SIZE) {
		uint8_t block[BLOCK_SIZE];
		sm4_crypt(ctx, BLOCK_SIZE, block, iv);
		for (size_t j = 0; j < BLOCK_SIZE && i + j < length; j++) {
			plaintext[i + j] = ciphertext[i + j] ^ block[j];
		}
		memcpy(iv, ciphertext + i, BLOCK_SIZE);
	}
}

// OFB 模式分段加密/解密（OFB 模式加解密操作相同）
static void sm4_ofb_encrypt_decrypt_segment(struct sm4_ctx *ctx, uint8_t *iv,
					    const uint8_t *input, size_t length,
					    uint8_t *output)
{
	// struct sm4_ctx ctx;
	// sm4_set_encrypt_key(&ctx, key);
	for (size_t i = 0; i < length; i += BLOCK_SIZE) {
		uint8_t block[BLOCK_SIZE];
		sm4_crypt(ctx, BLOCK_SIZE, block, iv);
		for (size_t j = 0; j < BLOCK_SIZE && i + j < length; j++) {
			output[i + j] = input[i + j] ^ block[j];
		}
		memcpy(iv, block, BLOCK_SIZE);
	}
}

struct nettle_sm4_data {
	struct sm4_ctx sm4_ctx;
	uint16_t alg;
};

static int init_ctx(void **ctx, uint8_t *key, uint32_t klen, uint8_t *iv,
		    uint16_t alg, bool isenc)
{
	struct nettle_sm4_data *data;

	data = malloc(sizeof(struct nettle_sm4_data));
	if (data == NULL)
		return -1;

	data->alg = alg;

	if (isenc)
		sm4_set_encrypt_key(&data->sm4_ctx, key);
	else
		sm4_set_decrypt_key(&data->sm4_ctx, key);

	*ctx = data;

	return 0;
}

static void free_ctx(void *ctx)
{
	free(ctx);

	return;
}

static int encrypt(void *ctx, uint8_t *iv, uint8_t *data, uint32_t len,
		   uint8_t *dst)
{
	struct nettle_sm4_data *sm4_data = ctx;

	switch (sm4_data->alg) {
	case ABCRYPTO_CIPHER_SM4_ECB:
		sm4_ecb_encrypt_segment(&sm4_data->sm4_ctx, data, len, dst);
		break;
	case ABCRYPTO_CIPHER_SM4_CBC:
		sm4_cbc_encrypt_segment(&sm4_data->sm4_ctx, iv, data, len, dst);
		break;
	case ABCRYPTO_CIPHER_SM4_CTR:
		sm4_ctr_encrypt_decrypt_segment(&sm4_data->sm4_ctx, iv, data,
						len, dst);
		break;
	case ABCRYPTO_CIPHER_SM4_CFB:
		sm4_cfb_encrypt_segment(&sm4_data->sm4_ctx, iv, data, len, dst);
		break;
	case ABCRYPTO_CIPHER_SM4_OFB:
		sm4_ofb_encrypt_decrypt_segment(&sm4_data->sm4_ctx, iv, data,
						len, dst);
		break;
	default:
		return -1;
	}

	return 0;
}

static int decrypt(void *ctx, uint8_t *iv, uint8_t *data, uint32_t len,
		   uint8_t *dst)
{
	struct nettle_sm4_data *sm4_data = ctx;

	switch (sm4_data->alg) {
	case ABCRYPTO_CIPHER_SM4_ECB:
		sm4_ecb_decrypt_segment(&sm4_data->sm4_ctx, data, len, dst);
		break;
	case ABCRYPTO_CIPHER_SM4_CBC:
		sm4_cbc_decrypt_segment(&sm4_data->sm4_ctx, iv, data, len, dst);
		break;
	case ABCRYPTO_CIPHER_SM4_CTR:
		sm4_ctr_encrypt_decrypt_segment(&sm4_data->sm4_ctx, iv, data,
						len, dst);
		break;
	case ABCRYPTO_CIPHER_SM4_CFB:
		sm4_cfb_decrypt_segment(&sm4_data->sm4_ctx, iv, data, len, dst);
		break;
	case ABCRYPTO_CIPHER_SM4_OFB:
		sm4_ofb_encrypt_decrypt_segment(&sm4_data->sm4_ctx, iv, data,
						len, dst);
		break;
	default:
		return -1;
	}

	return 0;
}

static int once_encrypt(uint8_t *key, uint8_t *iv, uint8_t *src, uint8_t *dst,
			uint32_t len, uint16_t alg)
{
	void *ctx = NULL;
	int ret;

	ret = init_ctx(&ctx, key, 16, iv, alg, true);
	if (ret)
		return -1;

	ret = encrypt(ctx, iv, src, len, dst);

	free_ctx(ctx);

	return ret;
}

static int once_decrypt(uint8_t *key, uint8_t *iv, uint8_t *src, uint8_t *dst,
			uint32_t len, uint16_t alg)
{
	void *ctx = NULL;
	int ret;

	ret = init_ctx(&ctx, key, 16, iv, alg, false);
	if (ret)
		return -1;

	ret = encrypt(ctx, iv, src, len, dst);

	free_ctx(ctx);

	return ret;
}

static struct abcrypto_skcipher _nettle_cipher = { .init_ctx = init_ctx,
						   .free_ctx = free_ctx,
						   .encrypt = encrypt,
						   .decrypt = decrypt,
						   .once_encrypt = once_encrypt,
						   .once_decrypt =
							   once_decrypt };

struct abcrypto_skcipher *get_nettle_cipher(void)
{
	return &_nettle_cipher;
}
